from subprocess import Popen,PIPE
import re

def ssd_sn_life():
  disk_count=1
  j=1
  sub_cmd=''
  mid_cmd=[]

  cmd = "/usr/sbin/smartctl --scan | grep 'megaraid' | awk '{print $1,$2,$3}'"
  proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
  stdout, stderr = proc.communicate()
  disk_raid=stdout.strip()

  if len(disk_raid) == 0:
     cmd="/usr/sbin/smartctl --scan | grep -vE 'megaraid|nvme' | awk '{print $1}'"
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     disk_jbod=stdout.strip()
     all_disk=disk_raid+disk_jbod
  else:
     all_disk=disk_raid

  test_disk= "*".join(all_disk.split("\n"))
  test_disk = test_disk+'*'
  #print "test_disk length:{}".format(len(test_disk))
  #print "=====ssdtest_disk:{}".format(test_disk)
  for item in test_disk[:]:
      if item is not '*' and j <= len(test_disk):
          sub_cmd=sub_cmd+item
          j+=1
      else:
          j+=1
          mid_cmd.append(sub_cmd)
          #print "sub_cmd:{}".format(sub_cmd)
          cmd= "/usr/sbin/smartctl {}  -i | grep -i 'Serial number' | awk '{{print $3}}'".format(sub_cmd)
          proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
          stdout, stderr = proc.communicate()
          sn=stdout.strip()
          #for key,value in ssdsn_disk.items():
          #    if sn==key:
                  #life,total_written=get_info(sn,sub_cmd)

          cmd="/usr/sbin/smartctl {} -i | grep 'SMART support is'".format(mid_cmd[-1])
          proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
          stdout, stderr = proc.communicate()
          smart_support=stdout.strip()
          if 'Available' in smart_support:
             disk_smart_status[sn]=1

          life,total_written,disk_type,capacity,model=get_life_written(mid_cmd)
          ssdsn_life[sn]=re.sub("\D","",life)
          ssdsn_written[sn]=total_written
          sn_type[sn]=disk_type
          sn_capacity[sn]=capacity
          sn_model[sn]=re.sub(r'\s+', ' ', model)
          #print "=====sn_type:{}".format(sn_type)
          #print "=====ssdsn_life:{}".format(ssdsn_life)
          #print "=====ssdsn_written:{}".format(ssdsn_written)
          reallocate_sector,reported_uncorrect,current_pend,cmd_timeout,offline_uncorrect=get_sata_info(mid_cmd)
          sata_reallocate[sn]=reallocate_sector
          sata_uncorrect[sn]=reported_uncorrect
          sata_pend[sn]=current_pend
          sata_timeout[sn]=cmd_timeout
          sata_offline_uncorrect[sn]=offline_uncorrect

          sub_cmd=''

def nvme_info():

    j=1
    sub_cmd=''
    mid_cmd=[]

    cmd = "/usr/sbin/smartctl --scan | grep 'nvme' | awk '{print $1,$2,$3}'"
    proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
    stdout, stderr = proc.communicate()
    disk=stdout.strip()

    test_disk= "*".join(disk.split("\n"))
    test_disk = test_disk+'*'

    #print test_life,type(test_life)
    for item in test_disk[:]:
        if item is not '*' and j <= len(test_disk):
            sub_cmd=sub_cmd+item
            j+=1
        else:
            j+=1
            mid_cmd.append(sub_cmd)
            cmd= "/usr/sbin/smartctl {}  -i | grep -i 'Serial number' | awk '{{print $NF}}'".format(sub_cmd)
            proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
            stdout, stderr = proc.communicate()
            sn=stdout.strip()

            disk=sub_cmd[5:10]
            nvmesn_disk[sn]=disk

            disk_smart_status[sn]=1
            life,total_written,disk_type,capacity,model=get_life_written(mid_cmd)
            nvmesn_life[sn]=re.sub("\D","",life)
            nvmesn_written[sn]=total_written
            sn_type[sn]=disk_type
            sn_capacity[sn]=capacity
            sn_model[sn]=re.sub(r'\s+', ' ', model)

            reallocate_sector,reported_uncorrect,current_pend,cmd_timeout,offline_uncorrect=get_sata_info(mid_cmd)
            sata_reallocate[sn]=reallocate_sector
            sata_uncorrect[sn]=reported_uncorrect
            sata_pend[sn]=current_pend
            sata_timeout[sn]=cmd_timeout
            sata_offline_uncorrect[sn]=offline_uncorrect
            sub_cmd=''


def get_life_written(mid_cmd):

     if 'nvme' in mid_cmd[-1]:
         cmd = "/usr/sbin/smartctl {} -x | grep 'Percentage Used' | awk '{{print $NF}}'".format(mid_cmd[-1])
     else:
         cmd = "/usr/sbin/smartctl {} -x | grep 'Percentage Used' | awk '{{print $4}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     life=stdout.strip()
     #sn_life[sn]=re.sub("\D","",life)

     cmd="/usr/sbin/smartctl {} -x | grep 'Logical Sectors Written' | awk '{{print $4}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     written=stdout.strip()
     #print "=====written:{}".format(written)

     cmd="/usr/sbin/smartctl {} -x | grep 'Sector Sizes' | awk '{{print $3}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     sector_size=stdout.strip()
     #print "=====sector_size:{}".format(sector_size)
     if written == '' or sector_size == '':
        total_written= -1
     else:
        total_written=int(written)*int(sector_size)
        #sn_written[sn]=total_written
    
     disk_type = ''
     capacity = ''
     cmd="/usr/sbin/smartctl {} -x | grep 'SMART support is'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     smart_support=stdout.strip()
     if 'Available' in smart_support or 'nvme' in mid_cmd[-1]:
          cmd="/usr/sbin/smartctl {} -x | grep 'Rotation Rate'".format(mid_cmd[-1])
          proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
          stdout, stderr = proc.communicate()
          temp_type=stdout.strip()
          if temp_type != '':
              if 'Solid State Device' in temp_type:
                  disk_type = 'SSD'
              elif 'rpm' in temp_type:
                  disk_type = 'HDD'
          
          if 'nvme' in mid_cmd[-1]:
               disk_type = 'SSD'
               cmd="/usr/sbin/smartctl {} -x |grep -i 'Total.*Capacity'| grep -oP '\\[\\K[^\\]]+'".format(mid_cmd[-1])
               proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
	       stdout, stderr = proc.communicate()
	       capacity=stdout.strip()
               
               if capacity == '':
                  cmd="/usr/sbin/smartctl {} -x |grep -i 'Namespace.*Size/Capacity'| grep -oP '\\[\\K[^\\]]+'".format(mid_cmd[-1])
                  proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
                  stdout, stderr = proc.communicate()
                  capacity=stdout.strip()
          else :
               cmd="/usr/sbin/smartctl {} -x | grep 'User Capacity'| awk -F : '{{print $2}}' | grep -oP '\\[\\K[^\\]]+'".format(mid_cmd[-1])
               proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
               stdout, stderr = proc.communicate()
               capacity=stdout.strip()
         
     cmd="smartctl {} -i".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     mid_model=stdout.strip().replace("\n"," ")  

     if any(keyword in mid_model for keyword in ['Model Number', 'Vendor', 'Model Family', 'Device Model']):
     #if 'nvme|Model Number|Vendor|Model Family|Device Model' in mid_cmd[-1]:
         mid_parts=mid_model.split('=== START OF INFORMATION SECTION ===',1)
         parts=mid_parts[1].split(':',1)
         sub_model=parts[1].strip().split()[0]
         model='Seagate' if sub_model.startswith('ST') else sub_model
     else:
         model = ' '
    # if 'nvme' in mid_cmd[-1]:
    #     if 'Model Number' in mid_model:
    #         mid_parts=mid_model.split('=== START OF INFORMATION SECTION ===',1)
    #         parts=mid_parts[1].split(':',1)
    #         model=parts[1].strip().split()[0]  
    # else:
    #     if 'Vendor' in mid_model:
    #         mid_parts=mid_model.split('=== START OF INFORMATION SECTION ===',1)
    #         parts=mid_parts[1].split(':',1)
    #         model=parts[1].strip().split()[0]
    #         #model='Seagate' if sub_model.startswith('ST') else sub_model
    #     elif 'Model Family' in mid_model:
    #         mid_parts=mid_model.split('=== START OF INFORMATION SECTION ===',1)
    #         parts=mid_parts[1].split(':',1)
    #         model=parts[1].strip().split()[0]
    #         #model='Seagate' if sub_model.startswith('ST') else sub_model
    #     elif 'Device Model' in mid_model:
    #         mid_parts=mid_model.split('=== START OF INFORMATION SECTION ===',1)
    #         parts=mid_parts[1].split(':',1)
    #         #print '-----parts:{}'.format(parts)
    #         sub_model=parts[1].strip().split()[0]
    #         model='Seagate' if sub_model.startswith('ST') else sub_model
    #     else:
    #         model=' '  

     return life,total_written,disk_type,capacity,model

def get_sata_info(mid_cmd):

     cmd="/usr/sbin/smartctl {} -A | grep '5 *reallocate' -i | awk '{{print $NF}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     reallocate_sector=stdout.strip()

     cmd="/usr/sbin/smartctl {} -A | grep -i '187.*uncorrect' | awk '{{print $NF}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     reported_uncorrect=stdout.strip()

     cmd="/usr/sbin/smartctl {} -A | grep -i '197.*pend' | awk '{{print $NF}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     current_pend=stdout.strip()


     cmd="/usr/sbin/smartctl {} -A | grep -i '188.*timeout' | awk '{{print $NF}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     cmd_timeout=stdout.strip()


     cmd="/usr/sbin/smartctl {} -A | grep -i '198.*uncorrect' | awk '{{print $NF}}'".format(mid_cmd[-1])
     proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
     stdout, stderr = proc.communicate()
     offline_uncorrect=stdout.strip()

     return reallocate_sector,reported_uncorrect,current_pend,cmd_timeout,offline_uncorrect

def print_info(dict_info,string):
    if len(dict_info) !=0:
       if dict_info == ssdsn_written or dict_info == nvmesn_written:
          for key,value in dict_info.items():
             if value!='' and value != -1:
                 print 'smartctl_device_{}{{sn="{}",type="{}",capacity="{}",model="{}"}}{:.5e}'.format(string,key,sn_type[key],sn_capacity[key].replace(' ',''),sn_model[key],value)
       else:
          for key,value in dict_info.items():
             if value!='':
                 print 'smartctl_device_{}{{sn="{}",type="{}",capacity="{}",model="{}"}}{}'.format(string,key,sn_type[key],sn_capacity[key].replace(' ',''),sn_model[key],value)        

if __name__=='__main__':
  ssdsn_life={} # dict   key--->sn   value--->PUEI
  ssdsn_written={}#dict   key--->sn   value--->TBW,total written
  nvmesn_disk={}
  nvmesn_life={}
  nvmesn_written={}

  sn_type={} #ssd or hdd
  sn_capacity={}
  sn_model={}#intel,...  

  sata_reallocate={}
  sata_uncorrect={}
  sata_pend={}
  sata_timeout={}
  sata_offline_uncorrect={}

  disk_smart_status={} # key--->sn
  #sata ssd, value--->SMART support is
  #nvme ssd, value--->1

  cmd = "/usr/sbin/smartctl --scan | grep 'nvme' | awk '{print $1,$2,$3}'"
  proc = Popen(cmd, stdout=PIPE, stderr=PIPE, shell=True)
  stdout, stderr = proc.communicate()
  disk=stdout.strip()
  if 'nvme' in disk:
      nvme_info()
  ssd_sn_life()

  print "# HELP smartctl_device_smart_status Disk Smart Status"
  print "# TYPE smartctl_device_smart_status counter"
  print_info(disk_smart_status,'smart_status')

  print "# HELP smartctl_device_percentage_used Percentage Used Endurance Indicatord"
  print "# TYPE smartctl_device_percentage_used counter"
  print_info(ssdsn_life,'percentage_used')
  print_info(nvmesn_life,'percentage_used')

  print "# HELP smartctl_device_bytes_written Logical Sectors Written Transformer to Bytes"
  print "# TYPE smartctl_device_bytes_written counter"
  print_info(ssdsn_written,'bytes_written')
  print_info(nvmesn_written,'bytes_written') 

  print "# HELP smartctl_device_reallocated_sector_ct Reallocated_Sector_Ct"
  print "# TYPE smartctl_device_reallocated_sector_ct counter"
  print_info(sata_reallocate,'reallocated_sector_ct')

  print "# HELP smartctl_device_reported_uncorrect Reported_Uncorrect"
  print "# TYPE smartctl_device_reported_uncorrect counter"
  print_info(sata_uncorrect,'reported_uncorrect')

  print "# HELP smartctl_device_current_pending_sector Current_Pending_Sector"
  print "# TYPE smartctl_device_current_pending_sector counter"
  print_info(sata_pend,'current_pending_sector')

  print "# HELP smartctl_device_command_timeout Command Timeout"
  print "# TYPE smartctl_device_command_timeout counter"
  print_info(sata_timeout,'command_timeout')

  print "# HELP smartctl_device_offline_uncorrectable Offline Uncorrectable"
  print "# TYPE smartctl_device_offline_uncorrectable counter"
  print_info(sata_offline_uncorrect,'offline_uncorrectable')
#
